from PIL import ImageGrab, Image, ImageFilter
import cv2
import time
import numpy as np
from utils import get_app_book, get_date_str
import sys


def add_border(cv_img, border_size=1, color=(0, 0, 0)):
    return cv2.copyMakeBorder(cv_img, 0, 0, border_size, 0, cv2.BORDER_CONSTANT, value=color)


def process_info(wb):
    info_sheet = wb.Sheets("info")
    if info_sheet is None:
        print("未找到info工作表。")
        return None
    
    # info sheet定义了该wb的基本信息，
    # 其中type可选值为：dr(日报)，bt(节拍)，an(异常)，vs(车型统计)，pl(计划),ot(其他)
    # start和end表示该sheet的范围(已弃用)
    # father_plan仅对type为dr类型有效，表示该日报所属的计划
    # date_row，step_col仅对type为pl类型有效，表示日期所在行和工序所在列(已弃用)
    col_idx = {
        "name":None,
        "type":None,
        "father_plan":None
    }

    data = info_sheet.UsedRange.Value
    # 获取表头
    headers = data[0]

    # 获取列索引
    for col, value in enumerate(headers):
        if value in col_idx:
            col_idx[value] = col
    sheets = []

    # 获取数据
    for row in data[1:]:
        sheet = {key: row[col_idx[key]] for key in col_idx if col_idx[key] is not None}
        sheets.append(sheet)
    return sheets


def capture_image(sheets_info, wb, debug=False):
    imgs = []
    for info in sheets_info:
        if True: # 后续开发可添加判断条件
            try:
                sheet = wb.Sheets(info["name"])
                sheet.UsedRange.CopyPicture(Appearance=1, Format=2)
                if debug:
                    print(f"成功复制图片: {info['name']}, 截图范围: {sheet.UsedRange.Address}")
                img = ImageGrab.grabclipboard()
                wait_time = 100
                # 等待获取到图像
                while img is None and wait_time > 0:
                    time.sleep(0.1)
                    wait_time -= 1
                
                if img is not None:
                    imgs.append(img)
                elif wait_time <= 0:
                    print(f"获取图像超时: {info['name']}")
            except Exception as e:
                print(f"获取图像失败: {e}")
    return imgs
            

def pil_to_cv2(pil_img, debug=False):
    """
    将 PIL.Image 转换为 OpenCV 格式(numpy 数组)，注意转换颜色通道顺序。
    """
    cv_img = np.array(pil_img)
    # 如果是 RGB 模式，则 Pillow 内部通道顺序为 RGB，而 OpenCV 使用 BGR
    if pil_img.mode == "RGB":
        if debug:
            print("RGB mode")
        cv_img = cv2.cvtColor(cv_img, cv2.COLOR_RGB2BGR)
    elif pil_img.mode == "RGBA":
        if debug:
            print("RGBA mode")
        cv_img = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGRA)
    elif debug:
        print("其他模式")
    
    if debug:
        print("转换成功")
    return cv_img


def zoom_image(cv_img, target_width=1090):
    """
    等比例缩放图片至指定宽度。
    """
    # 获取原始图像的宽度和高度
    h, w = cv_img.shape[:2]
    # 计算缩放比例
    ratio = target_width / w
    # 计算缩放后的尺寸
    dim = (target_width, int(h * ratio))

    # 若放大则使用 INTER_LANCZOS4，缩小则使用 INTER_AREA
    interp = cv2.INTER_LANCZOS4 if ratio > 1 else cv2.INTER_AREA
    # 缩放图像并返回
    return cv2.resize(cv_img, dim, interpolation=interp)


def zoom_pil_image(pil_img, target_width=1090):
    """
    等比例缩放图片至指定宽度。
    """
    # 获取原始图像的宽度和高度
    w, h = pil_img.size
    # 计算缩放比例
    ratio = target_width / w
    # 计算缩放后的尺寸
    dim = (target_width, int(h * ratio))

    # 缩放图像并返回
    return pil_img.resize(dim)


def split_image(cv_imgs, gap=20):
    total_height = sum([ img.shape[0] for img in cv_imgs ]) + gap * (len(cv_imgs) - 1)
    # 创建空白图像
    result = np.full((total_height, cv_imgs[0].shape[1], 3), 255, dtype=np.uint8)
    # 拼接图像
    y = 0
    for img in cv_imgs:
        h, w = img.shape[:2]
        result[y:y+h, :w] = img
        y += h + gap
    return result


def split_pil_image(pil_imgs, gap=20):
    total_height = sum([ img.size[1] for img in pil_imgs ]) + gap * (len(pil_imgs) - 1)
    # 创建空白图像
    result = Image.new("RGB", (pil_imgs[0].size[0], total_height), "white")
    # 拼接图像
    y = 0
    for img in pil_imgs:
        h = img.size[1]
        result.paste(img, (0, y))
        y += h + gap
    return result


def main(date=None, debug=True):
    if date is None:
        file_name = f"{get_date_str()}.xlsx"
    else:
        file_name = f"{date}.xlsx"
    # 打开工作簿
    wb, excel_app = get_app_book(file_name, debug)

    if wb is None:
        print("打开工作簿失败。")
        return
    
    if debug:
        print(f"成功打开工作簿: {wb.Name}")

    # 获取info工作表
    sheets_info = process_info(wb)
    if sheets_info is None:
        print("未找到info工作表，请检查工作簿。")
        return
    
    if debug:
        print(f"成功获取info工作表: {sheets_info}")
    
    pil_imgs = capture_image(sheets_info, wb, debug)

    # 检查是否为空数组
    if len(pil_imgs) == 0:
        print("未获取到任何图像。")
        return
    
    if debug:
        print(f"成功获取到{len(pil_imgs)}张图像。")
    
    # 将PIL图像转换为OpenCV格式
    # cv_imgs = [add_border(pil_to_cv2(img, debug)) for img in pil_imgs]
    # cv_imgs = [pil_to_cv2(img) for img in pil_imgs]

    # 缩放图像
    # resized_imgs = [ zoom_image(img) for img in cv_imgs ]

    # 拼接图像
    # result = split_image(resized_imgs)

    # 将file_name中的.xlsx替换为.png
    # img_name = file_name.replace(".xlsx", ".png")
    # 保存图像
    # cv2.imwrite(img_name, result)

    zoomed_imgs = [zoom_pil_image(img) for img in pil_imgs]
    result = split_pil_image(zoomed_imgs, 10)
    result = result.filter(ImageFilter.SHARPEN)
    result.save(file_name.replace(".xlsx", ".png"), "PNG")


    wb.Close(SaveChanges=False)

    if not False:
        excel_app.Quit()


args = sys.argv
if len(args) >1:
    print(args)
    main(args[1])
else:
    main()